#include "oppc.hpp"
#include <iostream>
#include <vector>
#include <fstream>
#include <iomanip>
#include <cmath> // For sqrt, M_PI
#include <algorithm> // For std::min, std::max

#ifndef __SYNTHESIS__
#include <cstdlib> // For srand, rand
#include <ctime>   // For time
#endif

// 辅助函数：将 RGB 颜色写入 BMP 文件 (简化版本，仅支持 24-bit BMP)
void write_bmp(const std::string& filename, int width, int height, const ap_uint<24>* data) {
    std::ofstream ofs(filename, std::ios::binary);
    if (!ofs) {
        std::cerr << "Error: Could not open BMP file for writing." << std::endl;
        return;
    }

    // BMP Header (14 bytes)
    ofs.put('B'); ofs.put('M'); // Signature
    ap_uint<32> file_size = 14 + 40 + width * height * 3;
    ofs.write((char*)&file_size, 4);
    ap_uint<32> reserved = 0;
    ofs.write((char*)&reserved, 4);
    ap_uint<32> data_offset = 14 + 40;
    ofs.write((char*)&data_offset, 4);

    // DIB Header (40 bytes)
    ap_uint<32> dib_header_size = 40;
    ofs.write((char*)&dib_header_size, 4);
    ap_uint<32> image_width = width;
    ofs.write((char*)&image_width, 4);
    ap_uint<32> image_height = height;
    ofs.write((char*)&image_height, 4);
    ap_uint<16> planes = 1;
    ofs.write((char*)&planes, 2);
    ap_uint<16> bpp = 24; // 24 bits per pixel
    ofs.write((char*)&bpp, 2);
    ap_uint<32> compression = 0; // No compression
    ofs.write((char*)&compression, 4);
    ap_uint<32> image_size_bytes = width * height * 3; // Renamed to avoid conflict with `image_size` from header
    ofs.write((char*)&image_size_bytes, 4);
    ap_uint<32> x_ppm = 0; // Horizontal resolution (pixels per meter)
    ofs.write((char*)&x_ppm, 4);
    ap_uint<32> y_ppm = 0; // Vertical resolution (pixels per meter)
    ofs.write((char*)&y_ppm, 4);
    ap_uint<32> colors_used = 0;
    ofs.write((char*)&colors_used, 4);
    ap_uint<32> colors_important = 0;
    ofs.write((char*)&colors_important, 4);

    // Pixel Data
    for (int y = height - 1; y >= 0; --y) { // BMP stores rows bottom-up
        for (int x = 0; x < width; ++x) {
            ap_uint<24> pixel = data[y * width + x];
            ofs.put((char)pixel.range(7, 0));  // Blue
            ofs.put((char)pixel.range(15, 8)); // Green
            ofs.put((char)pixel.range(23, 16)); // Red
        }
    }
    ofs.close();
}

// 全局静态纹理数据，以模拟 DDR 中的纹理
static ap_uint<32> global_ddr_texture[TEX_WIDTH * TEX_HEIGHT]; 

// 定义颜色常量，用于简化纹理生成
const ap_uint<32> TEX_COLOR_RED = 0xFF0000;
const ap_uint<32> TEX_COLOR_BLUE = 0x0000FF;

int main() {
    std::cout << "Starting OPPC C-simulation..." << std::endl;

    // --- 1. 初始化配置参数 ---
    OppeConfig config;
    // 光源从右上方斜向下照射，产生明暗对比
    config.light_dir = { (fpal::fp_vec_unit_t)0.707f, (fpal::fp_vec_unit_t)-0.707f, (fpal::fp_vec_unit_t)0.0f };
    
    config.light_color = { (fpal::fp_color_channel_t)255, (fpal::fp_color_channel_t)255, (fpal::fp_color_channel_t)255 };

    // 增加环境光强度，避免完全黑暗
    // --- 修改此处：增加环境光强度 ---
    config.ambient_color = { (fpal::fp_color_channel_t)150, (fpal::fp_color_channel_t)150, (fpal::fp_color_channel_t)150 }; // 从 80 调高到 150

    config.tile_start_x = 0;
    config.tile_start_y = 0;

    // --- 2. 模拟纹理数据 (DDR 中的 2048x2048 纹理) ---
    // 纹理现在只存储两种纯色：红色和蓝色，作为物体的基础颜色
    // 0 号纹理是红色，1 号纹理是蓝色
    global_ddr_texture[0] = TEX_COLOR_RED;  // 用于球体
    global_ddr_texture[1] = TEX_COLOR_BLUE; // 用于立方体
    // 其他纹理空间可以随意填充，这里只是为了演示
    for(int i = 2; i < TEX_WIDTH * TEX_HEIGHT; ++i) { // 使用 TEX_WIDTH * TEX_HEIGHT
        global_ddr_texture[i] = 0x000000; // 默认黑色
    }
    ap_uint<32>* p_ddr_texture = global_ddr_texture;

    // --- 3. 模拟输入像素片段流 ---
    hls::stream<PixelFragment> fragment_stream("fragment_stream");

    // 定义物体参数
    // 球体: 红色
    const float SPHERE_CENTER_X = (float)TILE_WIDTH / 2.0f + 5.0f; // 稍微偏右，靠前
    const float SPHERE_CENTER_Y = (float)TILE_HEIGHT / 2.0f;
    const float SPHERE_CENTER_Z = 0.6f; // 更近 (Z值越小越近)
    const float SPHERE_RADIUS = 10.0f;
    const float SPHERE_COLOR_UV_U = 0.0f; // 对应纹理索引 0 (红色)
    const float SPHERE_COLOR_UV_V = 0.0f;

    // 立方体: 蓝色
    const float CUBE_CENTER_X = (float)TILE_WIDTH / 2.0f - 5.0f; // 稍微偏左，靠后
    const float CUBE_CENTER_Y = (float)TILE_HEIGHT / 2.0f;
    const float CUBE_CENTER_Z = 0.8f; // 更远 (Z值越大越远)
    const float CUBE_HALF_SIDE = 10.0f;
    // --- 修改此处：将 CUBE_COLOR_UV_U 设置为 1.5f / TEX_WIDTH ---
    // 这样采样时 (U * (TEX_WIDTH - 1)) 的结果会在 [1.0, 2.0) 之间，取整后得到纹理索引 1
    const float CUBE_COLOR_UV_U = (1.0f + 0.5f) / (float)TEX_WIDTH; 
    const float CUBE_COLOR_UV_V = 0.0f;
    
    // Z-Buffer 深度映射：我们将 Z 坐标从 [0, 1] 映射到 ap_uint<16> 的 [0, 65535]
    // 0 是最近，1 是最远
    auto map_z_to_fixed = [](float z_val) {
        // Clamp z_val to [0.0, 1.0] to prevent out-of-range issues
        z_val = std::max(0.0f, std::min(1.0f, z_val));
        // 0.0f (最近) -> 0
        // 1.0f (最远) -> 65535 (对应 fpal::fp_z_buffer_t 的最大值)
        return (fpal::fp_z_buffer_t)(z_val * 65535.0f);
    };

    // 遍历 Tile 中的所有像素，**始终写入一个片段**
    for (int y = 0; y < TILE_HEIGHT; ++y) {
        for (int x = 0; x < TILE_WIDTH; ++x) {
            PixelFragment frag;
            frag.screen_x = config.tile_start_x + x;
            frag.screen_y = config.tile_start_y + y;
            // frag.is_valid_pixel = false; // PixelFragment 中已无此成员

            float closest_z = 1.0f; // 越小越近，初始化为最远
            bool pixel_covered = false; // 标记像素是否被任何物体覆盖
            fpal::fp_vec_unit_t final_normal_x = 0.0f;
            fpal::fp_vec_unit_t final_normal_y = 0.0f;
            fpal::fp_vec_unit_t final_normal_z = 0.0f;
            fpal::fp_vec_unit_t final_uv_u = 0.0f;
            fpal::fp_vec_unit_t final_uv_v = 0.0f;

            // --- 检查球体 ---
            float dist_sq_sphere = (x - SPHERE_CENTER_X) * (x - SPHERE_CENTER_X) +
                                   (y - SPHERE_CENTER_Y) * (y - SPHERE_CENTER_Y);
            if (dist_sq_sphere <= SPHERE_RADIUS * SPHERE_RADIUS) {
                float current_z = SPHERE_CENTER_Z;
                
                float normal_x_f = (x - SPHERE_CENTER_X);
                float normal_y_f = (y - SPHERE_CENTER_Y);
                float normal_z_f_sq = SPHERE_RADIUS * SPHERE_RADIUS - dist_sq_sphere;
                float normal_z_f = (normal_z_f_sq > 0.0f) ? std::sqrt(normal_z_f_sq) : 0.0f; 
                
                float normal_len = std::sqrt(normal_x_f*normal_x_f + normal_y_f*normal_y_f + normal_z_f*normal_z_f);
                if (normal_len > 0.001f) {
                    normal_x_f /= normal_len;
                    normal_y_f /= normal_len;
                    normal_z_f /= normal_len;
                } else {
                    normal_x_f = 0.0f; normal_y_f = 0.0f; normal_z_f = 1.0f; 
                }
                
                if (current_z < closest_z) {
                    closest_z = current_z;
                    final_normal_x = (fpal::fp_vec_unit_t)normal_x_f;
                    final_normal_y = (fpal::fp_vec_unit_t)normal_y_f;
                    final_normal_z = (fpal::fp_vec_unit_t)normal_z_f;
                    final_uv_u = (fpal::fp_vec_unit_t)SPHERE_COLOR_UV_U;
                    final_uv_v = (fpal::fp_vec_unit_t)SPHERE_COLOR_UV_V;
                    pixel_covered = true;
                }
            }

            // --- 检查立方体 ---
            float cube_min_x = CUBE_CENTER_X - CUBE_HALF_SIDE;
            float cube_max_x = CUBE_CENTER_X + CUBE_HALF_SIDE;
            float cube_min_y = CUBE_CENTER_Y - CUBE_HALF_SIDE;
            float cube_max_y = CUBE_CENTER_Y + CUBE_HALF_SIDE;

            if (x >= cube_min_x && x < cube_max_x &&
                y >= cube_min_y && y < cube_max_y) {
                float current_z = CUBE_CENTER_Z;
                fpal::vec3_fp_vec_unit cube_normal = {(fpal::fp_vec_unit_t)0.0f, (fpal::fp_vec_unit_t)0.5f, (fpal::fp_vec_unit_t)0.866f}; 

                if (current_z < closest_z) {
                    closest_z = current_z;
                    final_normal_x = cube_normal.x;
                    final_normal_y = cube_normal.y;
                    final_normal_z = cube_normal.z;
                    final_uv_u = (fpal::fp_vec_unit_t)CUBE_COLOR_UV_U; // 使用修改后的 CUBE_COLOR_UV_U
                    final_uv_v = (fpal::fp_vec_unit_t)CUBE_COLOR_UV_V;
                    pixel_covered = true;
                }
            }
            
            // 总是写入一个片段
            frag.interpolated_z = map_z_to_fixed(closest_z);
            frag.interpolated_uv.u = final_uv_u;
            frag.interpolated_uv.v = final_uv_v;
            frag.interpolated_normal = {final_normal_x, final_normal_y, final_normal_z};
            
            // 如果像素未被覆盖，提供一个默认的“背景”片段，其Z值为最远
            // oppc_core 的 z-buffer 初始化为 0xFFFF，所以最远的 Z 值应该能被正确处理
            if (!pixel_covered) {
                 frag.interpolated_z = map_z_to_fixed(1.0f); // 最远的 Z 值
                 // 也可以设置默认的 UV 或法线，虽然对于被 Z-Test 剔除的像素这些值无关紧要
                 frag.interpolated_uv = {(fpal::fp_vec_unit_t)0.0f, (fpal::fp_vec_unit_t)0.0f};
                 frag.interpolated_normal = {(fpal::fp_vec_unit_t)0.0f, (fpal::fp_vec_unit_t)0.0f, (fpal::fp_vec_unit_t)1.0f};
            }
            
            fragment_stream.write(frag);
        }
    }

    // --- 4. 准备输出 BRAM 接口 ---
    ap_uint<24> tile_framebuffer_output[TILE_PIXELS];
    fpal::fp_z_buffer_t tile_zbuffer_output[TILE_PIXELS];
    hls::stream<bool> tile_processed_signal("tile_processed_signal");

    // --- 5. 调用 OPPC 核心函数 ---
    oppc_core(
        fragment_stream,
        config,
        p_ddr_texture, // 提供模拟的 DDR 纹理数据指针
        tile_framebuffer_output, // 片上Framebuffer BRAM (作为输入/输出BRAM)
        tile_zbuffer_output,     // 片上Z-Buffer BRAM (作为输入/输出BRAM)
        tile_processed_signal
    );

    // --- 6. 验证结果 ---
    bool tile_done = false;
    // 使用阻塞式读取确保等待 oppc_core 完成。
    // 由于 oppc_core 保证写入一次 `true`，这里可以直接 `read()`。
    // 如果仿真仍然失败，可能需要检查 oppc_core 内部是否真的在写入 `true`。
    tile_processed_signal.read(tile_done); 
    
    if (tile_done) { // 检查读取到的值是否为 true
        std::cout << "OPPC module finished processing the tile successfully." << std::endl;
    } else {
        std::cerr << "Error: OPPC module did not signal tile completion with 'true'." << std::endl;
        return 1;
    }

    // 打印部分 Framebuffer 和 Z-Buffer 数据进行验证
    std::cout << "\nPartial Framebuffer Output (RGB):" << std::endl;
    for (int y = 0; y < TILE_HEIGHT; ++y) {
        for (int x = 0; x < TILE_WIDTH; ++x) {
            if (x < 8 && y < 8) { // 只打印左上角一小部分
                std::cout << std::hex << std::setfill('0') << std::setw(6) << tile_framebuffer_output[y * TILE_WIDTH + x] << " ";
            }
        }
        if (y < 8) std::cout << std::endl;
    }
    std::cout << std::dec << std::endl; // 恢复十进制输出

    std::cout << "\nPartial Z-Buffer Output:" << std::endl;
    for (int y = 0; y < TILE_HEIGHT; ++y) {
        for (int x = 0; x < TILE_WIDTH; ++x) {
            if (x < 8 && y < 8) { // 只打印左上角一小部分
                std::cout << std::setfill('0') << std::setw(5) << tile_zbuffer_output[y * TILE_WIDTH + x].to_uint() << " "; // 转换为 uint 方便显示
            }
        }
        if (y < 8) std::cout << std::endl;
    }

    // 将 Framebuffer 结果写入 BMP 文件进行可视化
    write_bmp("oppc_output_tile.bmp", TILE_WIDTH, TILE_HEIGHT, tile_framebuffer_output);
    std::cout << "\nRendered tile saved to oppc_output_tile.bmp" << std::endl;

    std::cout << "OPPC C-simulation finished." << std::endl;
    return 0;
}
